package com.zf.multilang.security;

import org.jsoup.Jsoup;
import org.jsoup.safety.Safelist;

import java.util.*;
import java.util.stream.Collectors;

public class XssUtils {

    // 白名单策略：仅允许基本标签（可根据业务调整）
    private static final Safelist BASIC_WITH_IMAGES = Safelist.basicWithImages()
        .addTags("h1", "h2", "h3", "h4", "h5", "h6", "p", "pre", "code", "blockquote")
        .addAttributes("img", "alt", "title")
        .addAttributes("a", "title");

    // 可配置：哪些字段允许 HTML（如富文本编辑器）
    private static final Set<String> ALLOWED_HTML_FIELDS = new HashSet<>(
        Arrays.asList("content", "description", "html", "editorData")
    );

    /**
     * 清理用户输入
     * @param input 输入字符串
     * @param fieldName 字段名（用于判断是否允许HTML）
     * @return 清理后的字符串
     */
    public static String clean(String input, String fieldName) {
        if (input == null) return null;

        // 保留换行符，Jsoup 会转成空格
        String preservedNewlines = input.replace("\n", "__XSS_NL__").replace("\r", "");

        String cleaned;
        if (isHtmlAllowed(fieldName)) {
            // 富文本字段：使用白名单过滤
            cleaned = Jsoup.clean(preservedNewlines, BASIC_WITH_IMAGES);
        } else {
            // 普通字段：完全转义
            cleaned = Jsoup.clean(preservedNewlines, Safelist.none());
        }

        return cleaned.replace("__XSS_NL__", "\n");
    }

    private static boolean isHtmlAllowed(String fieldName) {
        if (fieldName == null) return false;
        return ALLOWED_HTML_FIELDS.contains(fieldName.toLowerCase());
    }

    /**
     * 批量清理 Map（用于 JSON）
     */
    public static Map<String, Object> cleanMap(Map<String, Object> map) {
        Map<String, Object> result = new LinkedHashMap<>();
        for (Map.Entry<String, Object> entry : map.entrySet()) {
            String key = entry.getKey();
            Object value = entry.getValue();
            result.put(key, cleanRecursive(value, key));
        }
        return result;
    }

    @SuppressWarnings("unchecked")
    private static Object cleanRecursive(Object value, String fieldName) {
        if (value == null) {
            return null;
        } else if (value instanceof String) {
            return clean((String) value, fieldName);
        } else if (value instanceof Map) {
            return cleanMap((Map<String, Object>) value);
        } else if (value instanceof Collection) {
            return ((Collection<?>) value).stream()
                .map(item -> {
                    if (item instanceof Map) {
                        return cleanMap((Map<String, Object>) item);
                    } else if (item instanceof String) {
                        return clean((String) item, fieldName);
                    } else {
                        return item;
                    }
                }).collect(Collectors.toList());
        } else {
            return value;
        }
    }
}